# 混合搜索（语义 + 关键词）
from sklearn.feature_extraction.text import TfidfVectorizer
import numpy as np


class HybridRetriever:
    def __init__(self, embedding_model):
        self.embedding_model = embedding_model
        self.tfidf_vectorizer = TfidfVectorizer()
        self.tfidf_vectors = None

    def hybrid_retrieve(self, query: str, documents: List[str], alpha: float = 0.7) -> List[Dict]:
        """混合检索"""
        # 语义相似度
        query_embedding = self.embedding_model.encode([query])
        doc_embeddings = self.embedding_model.encode(documents)

        semantic_scores = np.dot(doc_embeddings, query_embedding.T).flatten()

        # TF-IDF 相似度
        if self.tfidf_vectors is None:
            self.tfidf_vectors = self.tfidf_vectorizer.fit_transform(documents)

        query_tfidf = self.tfidf_vectorizer.transform([query])
        keyword_scores = (
            query_tfidf @ self.tfidf_vectors.T).toarray().flatten()

        # 混合分数
        combined_scores = alpha * semantic_scores + \
            (1 - alpha) * keyword_scores

        # 返回排序结果
        results = []
        for i, score in enumerate(combined_scores):
            results.append({
                'content': documents[i],
                'semantic_score': float(semantic_scores[i]),
                'keyword_score': float(keyword_scores[i]),
                'combined_score': float(score)
            })

        results.sort(key=lambda x: x['combined_score'], reverse=True)
        return results
